# Written by Bram Cohen
# see LICENSE.txt for license information

import time ####### GROUP_VOD_CODE #######
from random import randrange, shuffle
from BitTornado.clock import clock
try:
    True
except:
    True = 1
    False = 0

class PiecePicker:
    def __init__(self, numpieces,
                 rarest_first_cutoff = 1, rarest_first_priority_cutoff = 3,
                 priority_step = 20):
        self.rarest_first_cutoff = rarest_first_cutoff
        self.rarest_first_priority_cutoff = rarest_first_priority_cutoff + priority_step
        self.priority_step = priority_step
        self.cutoff = rarest_first_priority_cutoff
        self.numpieces = numpieces
        self.started = []
        self.totalcount = 0
        self.numhaves = [0] * numpieces
        self.priority = [1] * numpieces
        self.removed_partials = {}
        self.crosscount = [numpieces]
        self.crosscount2 = [numpieces]
        self.has = [0] * numpieces
        self.numgot = 0
        self.done = False
        self.seed_connections = {}
        self.past_ips = {}
        self.seed_time = None
        self.superseed = False
        self.seeds_connected = 0
        self._init_interests()
        #self.to_print = None
        self.next_prefetch_chunk = None

    def _init_interests(self):
        self.interests = [[] for x in xrange(self.priority_step)]
        self.level_in_interests = [self.priority_step] * self.numpieces
        interests = range(self.numpieces)
        shuffle(interests)
        self.pos_in_interests = [0] * self.numpieces
        for i in xrange(self.numpieces):
            self.pos_in_interests[interests[i]] = i
        self.interests.append(interests)
        
    ####### GROUP_VOD_CODE #######
    def get_num_of_pieces(self):
        return self.numpieces    
    
    def computeWindowsSize(self):
        t = int(time.time() - self.streamwatcher.startTime)
        current_playing_point = int(((t - self.streamwatcher.delay) * self.streamwatcher.rate) / self.streamwatcher.toKbytes(self.streamwatcher.piece_size))
        if current_playing_point < 0:
            current_playing_point = 0
        
        prefetch_size_in_chunks = int((float(self.streamwatcher.prefetchT)/100) * (self.streamwatcher.numOfPieces))        
        
        self.window_size = prefetch_size_in_chunks*2 # this value is hard coded because the scripts doesn't support additional parameters
         
        #get next prefetch size... no point of downloading before it.
        if self.next_prefetch_chunk is None:
            self.next_prefetch_chunk = prefetch_size_in_chunks #dont download first preftech... 
        while self.next_prefetch_chunk <= current_playing_point:
            self.next_prefetch_chunk += prefetch_size_in_chunks
        
        window_start = self.next_prefetch_chunk
        
        window_end = window_start + self.window_size
        
        if window_end > self.streamwatcher.numOfPieces:
            window_end = self.streamwatcher.numOfPieces
        
        return window_start, window_end
    ####### GROUP_VOD_CODE #######

    def got_have(self, piece):
        self.totalcount+=1
        numint = self.numhaves[piece]
        self.numhaves[piece] += 1
        self.crosscount[numint] -= 1
        if numint+1==len(self.crosscount):
            self.crosscount.append(0)
        self.crosscount[numint+1] += 1
        if not self.done:
            numintplus = numint+self.has[piece]
            self.crosscount2[numintplus] -= 1
            if numintplus+1 == len(self.crosscount2):
                self.crosscount2.append(0)
            self.crosscount2[numintplus+1] += 1
            numint = self.level_in_interests[piece]
            self.level_in_interests[piece] += 1
        if self.superseed:
            self.seed_got_haves[piece] += 1
            numint = self.level_in_interests[piece]
            self.level_in_interests[piece] += 1
        elif self.has[piece] or self.priority[piece] == -1:
            return
        if numint == len(self.interests) - 1:
            self.interests.append([])
        self._shift_over(piece, self.interests[numint], self.interests[numint + 1])

    def lost_have(self, piece):
        self.totalcount-=1
        numint = self.numhaves[piece]
        self.numhaves[piece] -= 1
        self.crosscount[numint] -= 1
        self.crosscount[numint-1] += 1
        if not self.done:
            numintplus = numint+self.has[piece]
            self.crosscount2[numintplus] -= 1
            self.crosscount2[numintplus-1] += 1
            numint = self.level_in_interests[piece]
            self.level_in_interests[piece] -= 1
        if self.superseed:
            numint = self.level_in_interests[piece]
            self.level_in_interests[piece] -= 1
        elif self.has[piece] or self.priority[piece] == -1:
            return
        self._shift_over(piece, self.interests[numint], self.interests[numint - 1])

    def _shift_over(self, piece, l1, l2):
        assert self.superseed or (not self.has[piece] and self.priority[piece] >= 0)
        parray = self.pos_in_interests
        p = parray[piece]
        assert l1[p] == piece
        q = l1[-1]
        l1[p] = q
        parray[q] = p
        del l1[-1]
        #p = position of piece in l1
        #q = the last piece in l1
        #override p with q and delete p from l1 => piece is not in l1 anymore. parray[q]=parray[piece]
        newp = randrange(len(l2)+1)
        if newp == len(l2):
            parray[piece] = len(l2)
            l2.append(piece)
        else:
            old = l2[newp]
            parray[old] = len(l2)
            l2.append(old)
            l2[newp] = piece
            parray[piece] = newp


    def got_seed(self):
        self.seeds_connected += 1
        self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)

    def became_seed(self):
        self.got_seed()
        self.totalcount -= self.numpieces
        self.numhaves = [i-1 for i in self.numhaves]
        if self.superseed or not self.done:
            self.level_in_interests = [i-1 for i in self.level_in_interests]
            if self.interests:
                del self.interests[0]
        del self.crosscount[0]
        if not self.done:
            del self.crosscount2[0]

    def lost_seed(self):
        self.seeds_connected -= 1
        self.cutoff = max(self.rarest_first_priority_cutoff-self.seeds_connected,0)


    def requested(self, piece):
        if piece not in self.started:
            self.started.append(piece)

    def _remove_from_interests(self, piece, keep_partial = False):
        l = self.interests[self.level_in_interests[piece]]
        p = self.pos_in_interests[piece]
        assert l[p] == piece
        q = l[-1]
        l[p] = q
        self.pos_in_interests[q] = p
        del l[-1]
        try:
            self.started.remove(piece)
            if keep_partial:
                self.removed_partials[piece] = 1
        except ValueError:
            pass

    def complete(self, piece):

        ###### GROUP VOD ######
        # assert not self.has[piece]
        if (self.has[piece]):
            return
        #######################

        self.has[piece] = 1
        self.numgot += 1
        if self.numgot == self.numpieces:
            self.done = True
            self.crosscount2 = self.crosscount
        else:
            numhaves = self.numhaves[piece]
            self.crosscount2[numhaves] -= 1
            if numhaves+1 == len(self.crosscount2):
                self.crosscount2.append(0)
            self.crosscount2[numhaves+1] += 1
        self._remove_from_interests(piece)
        
    ####### GROUP_VOD_CODE #######
    ####### GROUP_VOD_CODE #######
    def next_in_order(self, begin, end, haves, wantfunc):
        for i in range(begin, end):
            dont_have_piece = (self.has[i]==0)
            didnt_request_it_all = wantfunc(i)
            the_other_client_have_it = haves[i]
            if dont_have_piece and didnt_request_it_all and the_other_client_have_it:
                #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(i)+'  from in_order' 
                return i
        return None
    
    def next_rarest_first(self, begin, end, haves, wantfunc, complete_first):
        cutoff = self.numgot < self.rarest_first_cutoff
        complete_first = (complete_first or cutoff) and not haves.complete()
        best = None
        bestnum = 2 ** 30
        for i in self.started:
            if i>=begin and i<end:
                if haves[i] and wantfunc(i):
                    if self.level_in_interests[i] < bestnum:
                        best = i
                        bestnum = self.level_in_interests[i]
        if best is not None:
            if complete_first or (cutoff and len(self.interests) > self.cutoff):
                #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(best)+'  from rarest_first - 1' 
                return best
        if haves.complete():
            r = [ (0, min(bestnum,len(self.interests))) ]
        elif cutoff and len(self.interests) > self.cutoff:
            r = [ (self.cutoff, min(bestnum,len(self.interests))),
                      (0, self.cutoff) ]
        else:
            r = [ (0, min(bestnum,len(self.interests))) ]
        for lo,hi in r:
            for i in xrange(lo,hi):
                for j in self.interests[i]:
                    if haves[j] and wantfunc(j) and j>=begin and j<end:
                        #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(j)+'  from rarest_first - 2' 
                        return j
        if best is not None:
            #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(best)+'  from rarest_first - 3' 
            return best
        return None
    
    def original_next(self, haves, wantfunc, complete_first):
        cutoff = self.numgot < self.rarest_first_cutoff
        complete_first = (complete_first or cutoff) and not haves.complete()
        best = None
        bestnum = 2 ** 30
        for i in self.started:
            if haves[i] and wantfunc(i):
                if self.level_in_interests[i] < bestnum:
                    best = i
                    bestnum = self.level_in_interests[i]
        if best is not None:
            if complete_first or (cutoff and len(self.interests) > self.cutoff):
                #self.to_print = 'piece number = '+str(best)+'  from original_next - 1' 
                return best
        if haves.complete():
            r = [ (0, min(bestnum,len(self.interests))) ]
        elif cutoff and len(self.interests) > self.cutoff:
            r = [ (self.cutoff, min(bestnum,len(self.interests))),
                      (0, self.cutoff) ]
        else:
            r = [ (0, min(bestnum,len(self.interests))) ]
        for lo,hi in r:
            for i in xrange(lo,hi):
                for j in self.interests[i]:
                    if haves[j] and wantfunc(j):
                        #self.to_print = 'piece number = '+str(j)+'  from original_next - 2' 
                        return j
        if best is not None:
            #self.to_print = 'piece number = '+str(best)+'  from original_next - 3' 
            return best
        return None
    
    def next_milestone2(self, haves, wantfunc, complete_first):
        current, end = self.computeWindowsSize()
        to_break = False
       
        #first, try in the immediate window 
        piece = self.next_in_order(current, end, haves, wantfunc)
        if piece is not None:
            return piece
        #try next windows, but rareset first on each window, also farest windows are bigger
        curr_iter = 1
        current = end
        end = end + self.window_size
        while (current<=self.numpieces and not to_break):
            if end>self.numpieces:
                end = self.numpieces
                to_break = True
            piece = self.next_rarest_first(current, end, haves, wantfunc, complete_first)
            if piece is not None:
                return piece
            current = end
            end = end+int(self.window_size * curr_iter)
            curr_iter += 1 #make window bigger each iter
        #can't find any good piece in the upcoming windows, just try rarest_first on all file
        return self.original_next(haves, wantfunc, complete_first)
    
    def next_just_in_order(self, haves, wantfunc, complete_first):
        current, end = self.computeWindowsSize()
        end = self.streamwatcher.numOfPieces # move the end to the end of the file so we will do in order on all the file
        piece = self.next_in_order(current, end, haves, wantfunc)
        if piece is None:
            return  self.original_next(haves, wantfunc, complete_first)
        else:
            return piece
    
    def next_rarest_first_in_order(self, begin, end, haves, wantfunc, complete_first):
        cutoff = self.numgot < self.rarest_first_cutoff
        complete_first = (complete_first or cutoff) and not haves.complete()
        have_best = False
        best = 2 ** 30
        bestnum = 2 ** 30
        for i in self.started:
            if i>=begin and i<end:
                if haves[i] and wantfunc(i):
                    if ((self.level_in_interests[i] < bestnum) or ((self.level_in_interests[i] == bestnum) and (i<best))):
                        best = i
                        bestnum = self.level_in_interests[i]
                        have_best = True
        if have_best:
            if complete_first or (cutoff and len(self.interests) > self.cutoff):
                #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(best)+'  from rarest_first - 1'
                return best
        if haves.complete():
            r = [ (0, min(bestnum,len(self.interests))) ]
        elif cutoff and len(self.interests) > self.cutoff:
            r = [ (self.cutoff, min(bestnum,len(self.interests))),
                      (0, self.cutoff) ]
        else:
            r = [ (0, min(bestnum,len(self.interests))) ]
        for lo,hi in r:
            for i in xrange(lo,hi):
                for j in self.interests[i]:
                    if haves[j] and wantfunc(j) and j>=begin and j<end:
                        #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(j)+'  from rarest_first - 2'
                        return j
        if have_best:
            #self.to_print = 'begin = '+str(begin)+'  end = '+str(end)+'  piece number = '+str(best)+'  from rarest_first - 3'
            return best
        return None
    
    def window_in_order_window_rarest_first(self, haves, wantfunc, complete_first):
        current, end = self.computeWindowsSize()
        #first, try in the immediate window 
        piece = self.next_in_order(current, end, haves, wantfunc)
        if piece is not None:
            return piece
        
        current = end
        end = self.streamwatcher.numOfPieces # move the end to the end of the file so we will do rarest first on the rest of the file
        ret = self.next_rarest_first_in_order (current, end, haves, wantfunc, complete_first)
        if ret is not None:
            return ret
        return self.original_next(haves, wantfunc, complete_first)
    
    def rarest_first_from_current(self, haves, wantfunc, complete_first):
        current, end = self.computeWindowsSize()
        end = self.streamwatcher.numOfPieces # move the end to the end of the file so we will do rarest first on the rest of the file
        ret = self.next_rarest_first(current, end, haves, wantfunc, complete_first)
        if ret is not None:
            return ret
        return self.original_next(haves, wantfunc, complete_first)

    #wantfunc is storageWrapper.do_I_have_requests which returns true if index is in inactiveRequets
    def next(self, haves, wantfunc, complete_first = False):
        ####### GROUP_VOD_CODE #######
        #return self.original_next(haves, wantfunc, complete_first)
        #return self.next_milestone2(haves, wantfunc, complete_first) #milestone 2
        #return self.window_in_order_window_rarest_first(haves, wantfunc, complete_first) #in order in one window and than find all the rarest first in the file and do in order on them
        #return self.rarest_first_from_current(haves, wantfunc, complete_first)
        #The Winning Algorithm
        return self.next_just_in_order(haves, wantfunc, complete_first) #in order on all the file

####### GROUP_VOD_CODE #######
####### GROUP_VOD_CODE #######

    def am_I_complete(self):
        return self.done
    
    def bump(self, piece):
        l = self.interests[self.level_in_interests[piece]]
        pos = self.pos_in_interests[piece]
        del l[pos]
        l.append(piece)
        for i in range(pos,len(l)):
            self.pos_in_interests[l[i]] = i
        try:
            self.started.remove(piece)
        except:
            pass

    def set_priority(self, piece, p):
        if self.superseed:
            return False    # don't muck with this if you're a superseed
        oldp = self.priority[piece]
        if oldp == p:
            return False
        self.priority[piece] = p
        if p == -1:
            # when setting priority -1,
            # make sure to cancel any downloads for this piece
            if not self.has[piece]:
                self._remove_from_interests(piece, True)
            return True
        if oldp == -1:
            level = self.numhaves[piece] + (self.priority_step * p)
            self.level_in_interests[piece] = level
            if self.has[piece]:
                return True
            while len(self.interests) < level+1:
                self.interests.append([])
            l2 = self.interests[level]
            parray = self.pos_in_interests
            newp = randrange(len(l2)+1)
            if newp == len(l2):
                parray[piece] = len(l2)
                l2.append(piece)
            else:
                old = l2[newp]
                parray[old] = len(l2)
                l2.append(old)
                l2[newp] = piece
                parray[piece] = newp
            if self.removed_partials.has_key(piece):
                del self.removed_partials[piece]
                self.started.append(piece)
            # now go to downloader and try requesting more
            return True
        numint = self.level_in_interests[piece]
        newint = numint + ((p - oldp) * self.priority_step)
        self.level_in_interests[piece] = newint
        if self.has[piece]:
            return False
        while len(self.interests) < newint+1:
            self.interests.append([])
        self._shift_over(piece, self.interests[numint], self.interests[newint])
        return False

    def is_blocked(self, piece):
        return self.priority[piece] < 0


    def set_superseed(self):
        assert self.done
        self.superseed = True
        self.seed_got_haves = [0] * self.numpieces
        self._init_interests()  # assume everyone is disconnected

    def next_have(self, connection, looser_upload):
        if self.seed_time is None:
            self.seed_time = clock()
            return None
        if clock() < self.seed_time+10:  # wait 10 seconds after seeing the first peers
            return None                 # to give time to grab have lists
        if not connection.upload.super_seeding:
            return None
        olddl = self.seed_connections.get(connection)
        if olddl is None:
            ip = connection.get_ip()
            olddl = self.past_ips.get(ip)
            if olddl is not None:                               # peer reconnected
                self.seed_connections[connection] = olddl
        if olddl is not None:
            if looser_upload:
                num = 1     # send a new have even if it hasn't spread that piece elsewhere
            else:
                num = 2
            if self.seed_got_haves[olddl] < num:
                return None
            if not connection.upload.was_ever_interested:   # it never downloaded it?
                connection.upload.skipped_count += 1
                if connection.upload.skipped_count >= 3:    # probably another stealthed seed
                    return -1                               # signal to close it
        for tier in self.interests:
            for piece in tier:
                if not connection.download.have[piece]:
                    seedint = self.level_in_interests[piece]
                    self.level_in_interests[piece] += 1  # tweak it up one, so you don't duplicate effort
                    if seedint == len(self.interests) - 1:
                        self.interests.append([])
                    self._shift_over(piece,
                                self.interests[seedint], self.interests[seedint + 1])
                    self.seed_got_haves[piece] = 0       # reset this
                    self.seed_connections[connection] = piece
                    connection.upload.seed_have_list.append(piece)
                    return piece
        return -1       # something screwy; terminate connection

    def lost_peer(self, connection):
        olddl = self.seed_connections.get(connection)
        if olddl is None:
            return
        del self.seed_connections[connection]
        self.past_ips[connection.get_ip()] = olddl
        if self.seed_got_haves[olddl] == 1:
            self.seed_got_haves[olddl] = 0
